#!/usr/bin/python
# encoding=utf-8
#
# swf.py - Quick and dirty .swf parser.
#
# Copyright 2007 National Library of Australia
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
Modified from: http://sourceforge.net/projects/swftool/
Modified by cosine
2010/1/6

SWF文件反编译工具
"""

from __future__ import generators
import os, sys, zlib, struct, math
import urllib2
from StringIO import *
from optparse import OptionParser

g_verbose = False
g_swfoutput = None # out decompiled swf file

# http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/142812
FILTER=''.join([(len(repr(chr(x)))==3) and chr(x) or '.' for x in range(256)])

def hexdump(src, length=8):
	N=0; result=''
	while src:
		s,src = src[:length],src[length:]
		hexa = ' '.join(["%02X"%ord(x) for x in s])
		s = s.translate(FILTER)
		result += "%04X   %-*s   %s\n" % (N, length*3, hexa, s)
		N+=length
	return result


TAGS = {12: 'DOACTION',
        26: 'PLACEOBJECT2',
        34: 'DEFINEBUTTON2',
        39: 'DEFINESPRITE',
        59: 'DOINITACTION' }


ACTIONS = {
    0x81: ("GotoFrame", "f"),
    0x83: ("GetUrl", "ul"),
    0x87: ("StoreRegister","r"),
    0x88: ("Constantpool", "Cc"),
    0x8a: ("WaitForFrame", "fs"),
    0x8b: ("SetTarget", "t"),
    0x8c: ("GotoLabel", "l"),
    0x8d: ("WaitForFrame2", "s"),
    0x94: ("With", "o"),
    0x96: ("Push", "p"),
    0x99: ("Jump", "b"),
    0x9a: ("GetUrl2", "m"),
    0x9b: ("DefineFunction", "{"),
    0x9d: ("If", "b"),
    0x9f: ("GotoFrame2", "m"),
}
"""
End of data.
"""

class StructIO:
	def __init__(self, f=None):
		if type(f) == str or f is None:
			f = StringIO(f)
		self._f = f
		self._leftover_bits = 8


		self.seek = f.seek
		self.tell = f.tell
		self.close = f.close
		self.write = f.write

		try:
			self.getvalue = f.getvalue
		except AttributeError:
			pass
	def __getstate__(self):
		odict = self.__dict__.copy()
		del odict['seek']
		del odict['tell']
		del odict['close']
		del odict['write']
		try:
			del odict['getvalue']
		except:
			pass
		return odict

	def __setstate__(self,dict):
		self.__dict__.update(dict)   # update attributes
		f = self._f

		self.seek = f.seek
		self.tell = f.tell
		self.close = f.close
		self.write = f.write

	def reads(self, fmt):
		size = struct.calcsize(fmt)
		data = self.read(size)
		if len(data) < size:
			raise EOFError
		return struct.unpack(fmt, data)

	def readstr(self):
		str = ''
		while 1:
			c = self.read(1)
			if c is '\0':
				break
			str += c
		return str


	def seek_byte(self):
		if self._leftover_bits != 8:
			# oh noes... we're partway through a byte
			# have to discard it
			self.seek(1, 1)
			self._leftover_bits = 8

	def read(self, bytes=None):
		self.seek_byte()
		if bytes is not None:
			return self._f.read(bytes)
		else:
			return self._f.read()

	def read_ub(self, size):
		x = 0
		read_bits = 0

		while read_bits < size:
			byte = ord(self._f.read(1))

			# read from our bit offset to the end
			# of this byte
			read_mask = (1 << self._leftover_bits) - 1
			val = byte & read_mask

			# if we've read too much 
			over = self._leftover_bits + read_bits - size
			if over > 0:
				# shift our value back to correct spot
				val >>= over
				read_this_time = self._leftover_bits - over

				# and restore the remaining bits
				self._leftover_bits = over
				self.seek(-1,1)
			else:
				# we read the whole byte, and didn't restore
				# any part of it, so we're sitting on a whole
				# new byte:
				read_this_time = self._leftover_bits
				self._leftover_bits = 8

			read_bits += read_this_time

			x <<= read_this_time
			x |= val
		return x

	def read_fixb(self, size):
		return self.read_ub(size) / 20.0

	def read_sb(self, size):
		x = self.read_ub(size)

		# if negative sign
		if x & (1<<(size-1)):
			# two's complement, flip the bits and add 1
			x = - (~x + 1) # fixme: untested

		return x


def reads(fmt, f):
	return f.reads(fmt)


class Block:
	def __init__(self, f, urllist):
		self.urllist = urllist
		start = f.tell()

		tag_head = reads("<H", f)[0]
		self.tagid = tag_head >> 6
		self.tag = TAGS.get(self.tagid, hex(self.tagid))        
		self.len = tag_head & 0x3f

		if self.len == 0x3f: # extended block
			self.len = reads("<I", f)[0]
			self.long = 1 # mark as long so as to have identical output to input
		else:
			self.long = 0

		self.data_s = f.read(self.len)
		self.data = StructIO(self.data_s)

		#print "{{", self.tagid, self.tag, self.len, "}}"
		#print hexdump(self.data_s, 16)

		self.test = f.tell() - start
	def __repr__(self):
		return '<%s (%d B)>' % (self.tag, self.len)

	def __str__(self):
		tag_head = self.tagid << 6
		size = len(self.data_s)
		if size < 0x3f and not self.long:
			tag_head |= size
			return struct.pack('<H', tag_head) + self.data_s
		else:
			tag_head |= 0x3f
			return struct.pack('<HI', tag_head, size) + self.data_s

	def do_local(self):
		pass

	def get_actions(self):
		return []

class Action:
	def __init__(self, opcode, f, urllist):
		self.opcode = opcode
		opname, self.argfmt = ACTIONS.get(opcode, (None, ''))

		if opcode < 0x80:
			length = 0
		else:
			length = f.reads('<H')[0]
		
		if g_verbose:
			print 'Action ', opname, '(%x)' % opcode, 'len', length

		self.args = []

		if not length:
			self.data = ''
			return

		start_act = f.tell()

		# read raw data in case we're unable to parse the arguments
		self.data = f.read(length)
		f.seek(start_act)

		#print 'Data:'
		#print hexdump(self.data)

		# attempt to parse args
		for arg in self.argfmt:
			if g_verbose:
				print '  (%s)  ' % arg,
			val = None
			if arg in ('u','l', 't'): # url, label, target
				val = f.readstr()
			elif arg == 'r': # register
				val = f.reads('<B')[0]
			elif arg == 'p': # push
				i = 0
				val = []
				while i < len(self.data):
					typecode = f.reads('<B')[0]
					if typecode == 0:  # string
						val.append( f.readstr() )
						i += len(val[-1]) + 2
					elif typecode == 1: # long
						val.append( f.reads('<I')[0])
						i += 5
					else:
						if g_verbose:
							print 'BAD TYPE', typecode
						val = None
						break
			elif arg == 'C': # constant pool count
				val = f.reads('<H')[0]
			elif arg == 'c': # constant pool values
				val = [f.readstr() for i in xrange(self.args[0])]
			if val is None:
				if g_verbose:
					print '???'
			else:
				if g_verbose:
					print val
				self.args.append(val)
				# add by cosine
				if arg == 'u' and val not in urllist:
					urllist.append(val)

		f.seek(start_act + length)
		if g_verbose:
			print

	def __str__(self):
		out = chr(self.opcode)

		if len(self.argfmt) == len(self.args) and self.argfmt:
			data = ''
			# arguments were parsed correctly, attempt to dump them
			for arg, val in zip(self.argfmt, self.args):
				if arg in ('u','l','t'):
					data += val + '\0'
				elif arg == 'r':
					data += chr(val)
				elif arg == 'p':
					for item in val:
						if isinstance(item, str):
							data += '\0' + item + '\0'
						elif isinstance(item, (int,long)):
							data += chr(1) + struct.pack('<I', item)
				elif arg == 'C': # constant pool count
					data += struct.pack('<H', val)
				elif arg == 'c': # constant pool strings
					for item in val:
						data += item + '\0'

		else:
			# just dump our raw data
			data = self.data 

		#if data != self.data:
		#        print 'Warning: Action modified: %s' % repr(self), len(self.argfmt), len(self.args)
		#        print 'before'
		#        print hexdump(data)
		#        print 'after'
		#        print hexdump(self.data)
		if self.opcode >= 0x80:
			out += struct.pack('<H', len(data)) + data
		return out 

	def __repr__(self):
		return ACTIONS.get(self.opcode, (hex(self.opcode),))[0]

class ActionCondition:
	def __init__(self, f, urllist):
		if g_verbose:
			print ' -> ActionCondition'
		self.cond = reads('<H', f)[0]

		self.actions = []
		# read actions
		while 1:
			opcode = f.reads('<B')[0]
			if not opcode: break

			self.actions.append(Action(opcode, f, urllist))
	def __str__(self):
		return struct.pack('<H', self.cond) + ''.join(map(str, self.actions)) + '\0'

class DoActionBlock(Block):
	def do_local(self):
		if self.tag == 'DOINITACTION':
			self.action_sprite = self.data.reads('<H')[0]

		# read actions
		self.actions = []
		while 1:
			opcode, = self.data.reads('<B')

			if not opcode: break
			self.actions.append(Action(opcode, self.data, self.urllist))
		if g_verbose:
			print self.actions

	def __str__(self):
		out = ''
		if self.tag == 'DOINITACTION':
			out += struct.pack('<H', self.action_sprite)
		out += ''.join(map(str, self.actions)) + '\x00'

		#if out != self.data_s:
		#    print 'DoAction mismatch!!'   
		#    print 'out.... ', len(out), repr(out)
		#    print 'data... ', len(self.data_s), repr(self.data_s)
		self.data_s = out
		return Block.__str__(self)


	def get_actions(self):
		return self.actions


class DefineButton2Block(Block):
	def do_local(self):

		self.button_id, self.flags = reads('<HB', self.data)

		self.data.reads('<H')

		if g_verbose:
			print '[Button]', self.button_id

		# read button records (we don't really care about the data)
		while self.data.reads('<B')[0]:
			char, layer = self.data.reads('<HH')

			Matrix(self.data)
			self.data.read(0) # realign
			CXForm(self.data)

		self.data.seek_byte()

		# mark where the conditions start for easy editing
		self.cond_offset = self.data.tell()

		self.conds = []

		# action conditions
		next = True
		try:
			while next:
				next = reads('<H', self.data)[0]
				self.conds.append(ActionCondition(self.data, self.urllist))
		except EOFError:
			sys.stderr.write('Warning: DefineButton2Block.do_local: abruptly reached end of block when reading actions.\n')

		if g_verbose:
			print '\n'


	def __str__(self):
		out = self.data_s[:self.cond_offset]
		for i,cond in zip(range(len(self.conds)), self.conds):
			tmp = str(cond)
			if i < len(self.conds)-1:
				next = len(tmp)+2 # fixme: is +2 correct??
			else:
				next = 0
			out += struct.pack('<H', next) + tmp
		self.data_s = out
		return Block.__str__(self)

	def get_actions(self):
		"Generator to retrieve all actions in the object."

		out = []
		for cond in self.conds:
			for action in cond.actions:
				out.append(action)
		return out

class PlaceObject2Block(Block):
	def do_local(self):
		global ParseVersion

		has_actions = self.data.read_ub(1)
		has_depth = self.data.read_ub(1)
		has_name = self.data.read_ub(1)
		has_ratio = self.data.read_ub(1)
		has_cxform = self.data.read_ub(1)
		has_matrix = self.data.read_ub(1)
		has_char = self.data.read_ub(1)
		flag_move = self.data.read_ub(1)

		depth = self.data.reads('<H')[0]
		if has_char:
			char = self.data.reads('<H')[0]
		if has_matrix:
			Matrix(self.data)
		if has_cxform:
			CXForm(self.data)
		if has_ratio:
			ratio = self.data.reads('<H')[0]
		if has_name:
			name = self.data.readstr()
		else:
			name = ''
		if has_depth:
			clip_depth = self.data.reads('<H')[0]

		self.action_offset = None
		if has_actions:
			# Oh, lovely, the object has actions associated with it
			if g_verbose:
				print '[Object]'

			# These are global flags, not sure what they do.
			flags = self.data.reads('<H')[0]
			if ParseVersion >= 6:
				more_flags = self.data.reads('<H')[0]

			# Not sure what this is... somehow it makes everything happy
			u1 = self.data.reads('<H')[0]

			self.action_offset = self.data.tell()
			self.conds = []

			while 1:
				# More flags. These ones define which event the
				# object is associated with.
				flags2 = self.data.reads('<H')[0]
				if ParseVersion >= 6:
					more_flags2 = self.data.reads('<H')[0]
				else:
					more_flags2 = None

				if not flags2 and not more_flags2: break

				cond = [(flags2, more_flags2)] 

				# Fixme: could be a 'keycode' byte here.

				size = self.data.reads('<I')[0]

				pos = self.data.tell()

				# Parse the actions
				while 1:
					opcode = self.data.reads('<B')[0]
					if not opcode: break
					cond.append( Action(opcode, self.data, self.urllist) )

				self.data.seek(pos + size)
				self.conds.append(cond)

			rest = self.data.read()
			if rest:
				if g_verbose:
					print repr(rest)
		#print '[Object]', name 

	def __str__(self):
		if self.action_offset is None:
			return Block.__str__(self)

		out = self.data_s[:self.action_offset]

		for i,cond in zip(range(len(self.conds)), self.conds):
			flags2, more_flags2 = cond[0]

			actions = ''.join(map(str, cond[1:])) + '\0'

			if more_flags2 is None:
				out += struct.pack('<HI', flags2, len(actions))
			else:
				out += struct.pack('<HHI', flags2, more_flags2, len(actions))

			out += actions


		# terminator
		if ParseVersion >= 6:
			out += struct.pack('<HH', 0,0)
		else:
			out += struct.pack('<H', 0)

		if self.data_s != out:
			if g_verbose:
				print 'XXXXX orig: ' + repr(self.data_s)
				print 'XXXXX  new: ' + repr(out)
		self.data_s = out
		return Block.__str__(self)

	def get_actions(self):
		"Generator to retrieve all actions in the object."

		if self.action_offset is None: return []

		out = []
		for cond in self.conds:
			for action in cond[1:]:
				out.append(action)
		return out

class DefineSpriteBlock(Block):
	"An embedded movie clip"
	def do_local(self):
		self.id = self.data.reads('<H')[0]
		self.blocks = []
		self.terminator = ''
		while self.data.tell() < len(self.data_s):
			# terminator doesn't seem to be consistant, sometimes
			# we have a trailing null other times we don't
			# so lets just copy whatever is there
			if  len(self.data_s) - self.data.tell() == 1:
				self.terminator = self.data.read(1)
				break

			self.blocks.append(block(self.data, self.urllist))
			#print 'n', self.data.tell(), len(self.data_s), repr(self.data_s[-1])
		if g_verbose:
			print self.blocks

	def __str__(self):
		out = struct.pack('<H', self.id)
		out += ''.join(map(str,self.blocks)) + self.terminator
		if out != self.data_s:
			if g_verbose:
				print 'Sprite mismatch!!'   
				print 'out.... ', len(out), repr(out)
				print 'data... ', len(self.data_s), repr(self.data_s)
		self.data_s = out
		return Block.__str__(self)



BLOCKS = {'DEFINEBUTTON2': DefineButton2Block, 
          'PLACEOBJECT2': PlaceObject2Block,
          'DEFINESPRITE': DefineSpriteBlock,
          'DOACTION': DoActionBlock,
          'DOINITACTION': DoActionBlock,
}
def block(f,urllist):
	'Block class selector'
	tag_head = reads("<H", f)[0]
	f.seek(-2, 1)
	tagid = tag_head >> 6
	tag = TAGS.get(tagid)

	blk = BLOCKS.get(tag, Block)

	x = blk(f,urllist)
	x.do_local()
	return x


class Matrix:
	'Matrix type'
	def __init__(self, f):
		has_scale = f.read_ub(1)
		if has_scale:
			scale_bits = f.read_ub(5)
			scalex = f.read_ub(scale_bits)
			scaley = f.read_ub(scale_bits)

		has_rot = f.read_ub(1)
		if has_rot:
			rot_bits = f.read_ub(5)
			rotx = f.read_ub(rot_bits)
			roty = f.read_ub(rot_bits)

		trans_bits = f.read_ub(5)
		transx = f.read_ub(trans_bits)
		transy = f.read_ub(trans_bits)

class CXForm:
	'Colour transform type'
	def __init__(self, f):
		has_mult = f.read_ub(1)
		has_add = f.read_ub(1)

		bits = f.read_ub(4)

		if has_mult:
			foo = f.read_ub(bits)
			foo = f.read_ub(bits)
			foo = f.read_ub(bits)

		if has_add:
			foo = f.read_ub(bits)
			foo = f.read_ub(bits)
			foo = f.read_ub(bits)

class Rect: 
	def __init__(self, f):


		bits = ord(f.read(1)) >> 3
		if g_verbose:
			print 'BITS', bits
		bytes = int(math.ceil((bits * 4 - 3 ) / 8.0))
		f.seek(-1, 1)
		self.raw_data = f.read(bytes+1)
		f.seek(-bytes-1, 1)

		bits = f.read_ub(5)
		if g_verbose:
			print 'BITS2', bits
		# fixme: should use a read sb instead
		self.xmin = f.read_fixb(bits)
		self.xmax = f.read_fixb(bits)
		self.ymin = f.read_fixb(bits)
		self.ymax = f.read_fixb(bits)

		if g_verbose:
			print self.xmin, self.xmax, self.ymin, self.ymax



	def __str__(self):
		# todo: actually restringify
		return self.raw_data

	def __repr__(self):
		return '<Rect %f, %f, %f, %f>' % (self.xmin, self.xmax, self.ymin, self.ymax)

class SWF:      
	def __init__(self, f):
		global ParseVersion
		# file header
		self.sig, self.ver, self.len = reads("<3sBI", f)
		ParseVersion = self.ver

		if self.sig == 'FWS':
			self.data_s = f.read()
		elif self.sig == 'CWS':
			self.data_s = zlib.decompress(f.read())
		else:
			raise ValueError, "File not in SWF format. (invalid signature: %s)" % repr(self.sig)

		self.data = StructIO(self.data_s)


		# frame header
		self.frame_size =  Rect(self.data) # movie size
		self.frame_rate, self.frame_count = reads("<HH", self.data)
		self.frame_rate /= float(2**8)

		if g_verbose:
			print self.frame_size, self.frame_rate

		self.blocks = []
		self.urllist = []
		try:
			while 1: self.blocks.append(block(self.data, self.urllist))
		except EOFError:
			pass

	def __str__(self):
		# todo: compression
		buf = StructIO()
		buf.write(struct.pack("<3sBI", 'FWS', self.ver, self.len))
		buf.write(str(self.frame_size))
		buf.write(struct.pack("<HH", int(self.frame_rate * 2**8), self.frame_count))

		for block in self.blocks:
			data = str(block)
			buf.write(data)

		return buf.getvalue()

	def iterstrings(self, blocks=None):
		if blocks is None: blocks = self.blocks
		for block in blocks:
			if hasattr(block, 'blocks'):
				for s in self.iterstrings(block.blocks):
					yield s

			for action in block.get_actions():
				if g_verbose:
					print repr(action)
				for i, arg in zip(range(len(action.args)), action.args):
					if isinstance(arg, str):
						yield arg, (action, i)
					elif isinstance(arg, list):
						for n, x in zip(range(len(arg)), arg):
							yield x, (action, i, n)

	def set_string(self, ref, s):
		s = str(s) # ensure non-unicode
		if len(ref) == 2:
			ref[0].args[ref[1]] = s
		elif len(ref) == 3:
			ref[0].args[ref[1]][ref[2]] = s


def srep_str(hs, s, r):
	if not isinstance(hs, str): return hs
	return hs.replace(s, r) 

def srep_block(block, s, r):
	"Run a string replacement regular expression over actions within a block"

	if hasattr(block, 'blocks'): [srep_block(x,s,r) for x in block.blocks]

	for action in block.get_actions():
		if g_verbose:
			print repr(action)
		for i, arg in zip(range(len(action.args)), action.args):
			if isinstance(arg, str):
				action.args[i] = srep_str(arg, s, r)
			elif isinstance(arg, list):
				action.args[i] = [srep_str(x,s,r) for x in arg]
			if g_verbose:
				print arg

def main():
	fn = sys.argv[1]
	f = file(fn, 'rb')
	x = SWF(StructIO(f))

	#for block in x.blocks:
	#    srep_block(block, 'SEARCH', 'REPLACE')

	f = file('out.swf', 'wb')
	f.write(str(x))
	return 0

def de_url(swfurl):
	"""return links"""
	links = []
	try:
		user_agent = 'Mozilla/4.0 (compatible; MSIE 8.0; Windows NT 6.0; Trident/4.0)'
		req = urllib2.Request(swfurl)
		req.add_header('User-Agent',user_agent)
		u = urllib2.urlopen(req)
		swfdata = u.read()
		u.close()
	except Exception,e:
		print 'read %s error: %s'%(swfurl,e)
		return links
	
	x = SWF(StructIO(swfdata))
	if g_swfoutput:
		f = file(g_swfoutput, 'wb')
		f.write(str(x))
		
	links = x.urllist
	return links
	
def de_file(swffile):
	"""return links"""
	links = []
	try:
		swfdata = file(swffile, 'rb')
	except Exception,e:
		print 'read %s error: %s'%(swffile,e)
		return links
	
	x = SWF(StructIO(swfdata))
	if g_swfoutput:
		f = file(g_swfoutput, 'wb')
		f.write(str(x))
		
	links = x.urllist
	return links
	
def de_data(swfdata):
	"""return links"""
	links = []
	x = SWF(StructIO(swfdata))
	links = x.urllist
	return links
	
if __name__ == '__main__':
	usage = "python swfparse.py -f foo.swf -v\n\
	python swfparse.py -u http://www.x.com/foo.swf -v"
	parser = OptionParser(usage)
	parser.add_option("-u", "--url", dest="swfurl",
	                  help="swf url like: http://www.x.com/foo.swf")
	parser.add_option("-f", "--file", dest="swffile",
		              help="local swf file")
	parser.add_option("-o", "--output", dest="output",
		              help="output decompiled swf file")
	parser.add_option("-v", action="store_true", dest="verbose",default=False,
		              help="print verbose info on screen")
	
	(options, args) = parser.parse_args()
	try:
		g_swfurl = options.swfurl
		g_swffile = options.swffile
		g_swfoutput = options.output
		g_verbose = options.verbose
	except:
		parser.print_help()
		sys.exit(0)
		
	if not g_swffile and not g_swfurl:
		parser.print_help()
		sys.exit(1)
	
	if g_swfurl:
		print '[*] decompile start: %s.'%g_swfurl
		print de_url(g_swfurl)
	if g_swffile:
		print '[*] decompile start: %s.'%g_swffile
		print de_file(g_swffile)
	
	print '[*] decompile end.'
	
